import json
import os
from datetime import datetime

import matplotlib.pyplot as plt
from matplotlib import font_manager
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# 设置中文字体
plt.rcParams['font.family'] = 'SimHei'  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 正确显示负号

# 设置训练集、验证集和测试集的路径
train_dir = 'yolo_dataset/train'
validation_dir = 'yolo_dataset/validation'
test_dir = 'yolo_dataset/test'

# 创建数据增强的 ImageDataGenerator 实例
train_datagen = ImageDataGenerator(
    rescale=1./255,
    rotation_range=20,
    width_shift_range=0.2,
    height_shift_range=0.2,
    shear_range=0.2,
    zoom_range=0.2,
    horizontal_flip=True,
    fill_mode='nearest'
)

validation_datagen = ImageDataGenerator(rescale=1./255)
test_datagen = ImageDataGenerator(rescale=1./255)

# 加载训练集
train_generator = train_datagen.flow_from_directory(
    train_dir,
    target_size=(224, 224),
    batch_size=10,  # 根据需要设置
    class_mode='categorical'
)
# 获取类别标签
class_labels = list(train_generator.class_indices.keys())
with open('model/class_labels.json', 'w', encoding='utf-8') as f:
    json.dump(class_labels, f)
print(class_labels)  # 输出类别名称

# 加载验证集
validation_generator = validation_datagen.flow_from_directory(
    validation_dir,
    target_size=(224, 224),
    batch_size=10,  # 根据需要设置
    class_mode='categorical'
)

# 加载测试集
test_generator = test_datagen.flow_from_directory(
    test_dir,
    target_size=(224, 224),
    batch_size=10,  # 根据需要设置
    class_mode='categorical',
    shuffle=False
)

# 检查样本数量
print(f"训练集样本数量: {train_generator.samples}")
print(f"验证集样本数量: {validation_generator.samples}")
print(f"测试集样本数量: {test_generator.samples}")

if (train_generator.samples < 1 or
    validation_generator.samples < 1 or
    test_generator.samples < 1):
    raise ValueError("训练集、验证集或测试集的样本数量太少，请检查数据集。")

# 加载预训练的 MobileNetV2 模型
base_model = MobileNetV2(weights='imagenet', include_top=False)

# 添加自定义的顶层分类器
x = base_model.output
x = GlobalAveragePooling2D()(x)
predictions = Dense(train_generator.num_classes, activation='softmax')(x)

# 构建完整的模型
model = Model(inputs=base_model.input, outputs=predictions)

# 冻结预训练模型的权重
for layer in base_model.layers:
    layer.trainable = False

# 编译模型
model.compile(optimizer=Adam(), loss='categorical_crossentropy', metrics=['accuracy'])

# 进行训练并输出验证结果
history = model.fit(
    train_generator,
    epochs=10,
    validation_data=validation_generator,
    validation_steps=validation_generator.samples // validation_generator.batch_size
)

# 评估模型在测试集上的表现
try:
    # 确保测试集样本数量大于0
    if test_generator.samples > 0:
        # 动态设置 steps
        steps = max(test_generator.samples // test_generator.batch_size, 1)
        test_loss, test_accuracy = model.evaluate(test_generator, steps=steps)
    else:
        raise ValueError("测试集没有样本。")
except Exception as e:
    print(f"测试集评估时出错: {e}")
    test_loss, test_accuracy = None, None

# 生成带时间戳的模型文件名
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
model_filename = f'model/mushroom_classification_model_{timestamp}.keras'
model.save(model_filename)

# 输出结果
print("训练过程中的结果：")
print("训练损失:", history.history['loss'])
print("训练准确率:", history.history['accuracy'])
print("验证损失:", history.history.get('val_loss', '无验证损失数据'))
print("验证准确率:", history.history.get('val_accuracy', '无验证准确率数据'))
if test_loss is not None and test_accuracy is not None:
    print(f"测试损失: {test_loss:.4f}")
    print(f"测试准确率: {test_accuracy:.4f}")

# 绘制图表
plt.figure(figsize=(12, 6))

# 绘制训练和验证准确率
plt.subplot(1, 2, 1)
plt.plot(history.history['accuracy'], label='训练准确率')
plt.plot(history.history.get('val_accuracy', []), label='验证准确率')
if test_accuracy is not None:
    plt.axhline(y=test_accuracy, color='r', linestyle='--', label='测试准确率')
plt.title('训练、验证和测试准确率')
plt.xlabel('轮次')
plt.ylabel('准确率')
plt.legend()

# 绘制训练和验证损失图表
plt.subplot(1, 2, 2)
plt.plot(history.history['loss'], label='训练损失')
plt.plot(history.history.get('val_loss', []), label='验证损失')
plt.title('训练和验证损失')
plt.xlabel('轮次')
plt.ylabel('损失')
plt.legend()

# 生成带时间戳的文件名
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
plot_filename = f'training_validation_plot_{timestamp}.png'

# 保存图表
plt.savefig(plot_filename)

# 显示图表
plt.tight_layout()
plt.show()

# 在训练完成后，保存模型信息
def save_model_info(model_path, history, test_accuracy=None):
    """保存模型信息到JSON文件"""
    model_info = {
        "model_path": model_path,
        "timestamp": datetime.now().strftime("%Y%m%d_%H%M%S"),
        "training_accuracy": float(history.history['accuracy'][-1]),
        "validation_accuracy": float(history.history['val_accuracy'][-1]) if 'val_accuracy' in history.history else None,
        "test_accuracy": float(test_accuracy) if test_accuracy is not None else None,
        "epochs": len(history.history['accuracy']),
        "training_history": {
            "accuracy": [float(x) for x in history.history['accuracy']],
            "loss": [float(x) for x in history.history['loss']],
            "val_accuracy": [float(x) for x in history.history['val_accuracy']] if 'val_accuracy' in history.history else None,
            "val_loss": [float(x) for x in history.history['val_loss']] if 'val_loss' in history.history else None
        }
    }
    
    # 确保model目录存在
    os.makedirs('model', exist_ok=True)
    
    # 读取现有的模型信息
    models_file = 'model/models_info.json'
    if os.path.exists(models_file):
        with open(models_file, 'r', encoding='utf-8') as f:
            models_list = json.load(f)
    else:
        models_list = []
    
    # 添加新的模型信息
    models_list.append(model_info)
    
    # 保存更新后的模型信息
    with open(models_file, 'w', encoding='utf-8') as f:
        json.dump(models_list, f, ensure_ascii=False, indent=2)
    
    return model_info

# 保存模型信息
model_info = save_model_info(
    model_path=model_filename,
    history=history,
    test_accuracy=test_accuracy
)
print("模型信息已保存到 model/models_info.json")
